using System;
using System.Collections.Generic;
using System.Threading;
using Cysharp.Threading.Tasks;

namespace MyUtility.UniTaskTool
{
      
    /// <summary>
    /// 带返回值得异步任务节点
    /// </summary>
    /// <typeparam name="T"></typeparam>
    public class UniTaskNode<T>
    {
        public static UniTaskNode<T> CreateTask(Func<CancellationToken, UniTask<T>> task, CancellationTokenSource cancellationTokenSource = default)
        {
            _id += 1;
            var taskNode = new UniTaskNode<T>(_id, task, cancellationTokenSource);
            _allNodes.TryAdd(_id, taskNode);
            return taskNode;
        }

        public static async UniTask<T> DoTask(Func<CancellationToken, UniTask<T>> task, CancellationTokenSource cancellationTokenSource = default)
        {
            _id += 1;
            var taskNode = new UniTaskNode<T>(_id, task, cancellationTokenSource);
            _allNodes.TryAdd(_id, taskNode);

            var val = await taskNode.DoTask();
            return val;
        }
        
        public static UniTaskNode<T> DoVoidTask(Func<CancellationToken, UniTask<T>> task, CancellationTokenSource cancellationTokenSource = default)
        {
            _id += 1;
            var taskNode = new UniTaskNode<T>(_id, task, cancellationTokenSource);
            _allNodes.TryAdd(_id, taskNode);

            taskNode.DoVoid();
            return taskNode;
        }
        
        
        public int Id { get; }
        public T Result { get; private set; }
        public CancellationToken Token => CancelTokenSrc.Token;
        public Func<CancellationToken, UniTask<T>> TaskFn { get; }

        public CancellationTokenSource CancelTokenSrc 
        {
            get
            {
                if (_cancelTokenSrc == null || _cancelTokenSrc.IsCancellationRequested)
                    _cancelTokenSrc = new CancellationTokenSource();
                return _cancelTokenSrc;
            }   
        }

        public UniTaskNode(int id, Func<CancellationToken, UniTask<T>> task)
        {
            Id = id;
            TaskFn = task;
            Result = default;
            _cancelTokenSrc = new CancellationTokenSource();
        }
        
        public UniTaskNode(int id, Func<CancellationToken, UniTask<T>> task, CancellationTokenSource cancelTokenSrc)
        {
            Id = id;
            TaskFn = task;
            Result = default;
            _cancelTokenSrc = cancelTokenSrc;
        }

        public void DoVoid()
        {
            DoTask().Forget();
        }

        public async UniTask<T> DoTask()
        {
            try
            {
                if (TaskFn == null)
                {
                    return default;
                }
                
                Result = await TaskFn.Invoke(Token);
                return Result;
            }
            finally
            {
                if(_allNodes.ContainsKey(Id)) _allNodes.Remove(Id);    
            }
        }

        public void CancelTask()
        {
            CancelTokenSrc.Cancel();
        }

        private CancellationTokenSource _cancelTokenSrc;

        static UniTaskNode()
        {
            _id = 0;
            _allNodes = new Dictionary<int, UniTaskNode<T>>();
        }

        private static int _id;
        private static Dictionary<int, UniTaskNode<T>> _allNodes;
        
    }
}